##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(ggplot2)
library(ggsci)
library(RColorBrewer)
library(forestplot)
library(patchwork)

##########################################################################################

option_list <- list(
    make_option(c("--input_burden_file"), type = "character"),
    make_option(c("--molecular_type"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    input_burden_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/molecular_type/logistic/baseInfo_mutBurden.tsv"
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/molecular_type/logistic"
}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

input_burden_file <- opt$input_burden_file
molecular_type <- opt$molecular_type
out_path <- opt$out_path

dir.create( out_path , recursive = T )

###########################################################################################

dat_burden <- data.frame(fread(input_burden_file))
dat_burden <- dat_burden[order(dat_burden$time , decreasing=T),]

###########################################################################################

trans <- function(num){
    text <- format(num,scientific=TRUE,digit=3)
    return(text)
}

###########################################################################################
## 包年、年龄
dat_smoke <- rbind( subset( dat_burden , baseType=="Tobacco" ))
dat_alcohol <- rbind( subset( dat_burden , baseType=="Alcohol" ))
dat_hp <- rbind( subset( dat_burden , baseType=="HP" ))

###########################################################################################
## 森林图
dat <- rbind(
    "" ,dat_smoke , 
    "" , dat_alcohol ,
    "" , dat_hp
    )

dat$p <- as.numeric(dat$p)
dat$p <- ifelse( dat$p < 0.01 , trans(dat$p) , round(dat$p , 2) )

col_risk <- c(
    rgb( red = 150 , green = 55 , blue = 53 , maxColorValue = 255 ) ,
    rgb( red = 127 , green = 127 , blue = 127 , maxColorValue = 255 ) ,
    rgb( red = 74 , green = 128 , blue = 189 , maxColorValue = 255 ) 
    )
names(col_risk) <- c("DGC" , "IGC" , "IM")

dat$b_clrs = ifelse( dat$time == "IM" , col_risk["IM"] , "" )
dat$b_clrs = ifelse( dat$time == "IGC" , col_risk["IGC"] , dat$b_clrs )
dat$b_clrs = ifelse( dat$time == "DGC" , col_risk["DGC"] , dat$b_clrs )

dat$baseUse <- c(
    "Smoking" , "" , "" , "" ,
    "Drinking" , "" , "" , "" ,
    "HP Status"  , "" , "" , ""
    )

dat$TimeUse <- c(
    "" , "IM" , "IGC" , "DGC" ,
    "" , "IM" , "IGC" , "DGC" ,
    "" , "IM" , "IGC" , "DGC"  
    )

iqr <- paste0( round(as.numeric(dat$coef) , 2) , " [" , 
    round(as.numeric(dat$coef_025) , 2) , "-" , 
    round(as.numeric(dat$coef_975) , 2) , "]"
    )
iqr[iqr=="NA [NA-NA]"] <- ""

tabletext <- cbind( 
    c("Population" , dat$baseUse) , 
    c("TMB" , dat$TimeUse) ,
    c("Coefficient" , iqr ) ,
    c("P" ,dat$p) )

cochrane_from_rmeta <- data.frame(
    mean  = c(NA, as.numeric(dat$coef)),
    lower = c(NA, as.numeric(dat$coef_025)),
    upper = c(NA, as.numeric(dat$coef_975)))

## 定义颜色
fn <- local({
  i = 0
  no_lines <- sum(!is.na(cochrane_from_rmeta$mean))
  b_clrs = dat$b_clrs
  b_clrs = b_clrs[b_clrs!=""]

  function(..., clr.line, clr.marker){
    i <<- i + 1
    fpDrawNormalCI(..., clr.marker = b_clrs[i])
  }
})

## 森林图
p <- forestplot(
        labeltext = tabletext , 
        cochrane_from_rmeta,
        # 添加水平线
        hrzl_lines = list("1" = gpar(lty=2, lwd=2, col="black"), 
                         "2" = gpar(lty=2, lwd=2, col="black")),
        align = "l", # 设置左边表格中字体的对齐方式
        zero = 0, # 设置zero line的位置
        new_page = TRUE,
        is.summary=c(
            TRUE,
            TRUE,
            rep(FALSE,3),
            TRUE,
            rep(FALSE,3),
            TRUE,
            rep(FALSE,3)
        ), #A vector indicating by TRUE/FALSE if the value is a summary value which means that it will have a different font-style
        #xticks = c(0.8, 1 , 1.5 , 2 , 2.5 , 3) ,
        boxsize = 0.25 ,
        lwd.ci = 2 , ## 横线的粗细
        col=fpColors(
            #box="#004C99",
            line="#004C99", 
            summary="royalblue", 
            hrz_lines = "#444444"),
        fn.ci_norm = fn,
        #xlog = TRUE ,
        graph.pos = 3 ,
        #xlog = TRUE , 
        #xticks = c(0.1, 1 , 2 , 3 , 4 , 5 , 7),
        vertices = TRUE
    )


out_file <- paste0(out_path , "/baseInfo_coef.",molecular_type,".pdf")
pdf(out_file , width = 6.52 , height = 3.81)
## Legend外部定义,par设置为0，保证legend可以在图的任意位置出现
plot.new()
par(mai=c(0,0,0,0))
p
#legend(x = -0.15, y = -0.4 ,  
#inset=.02, 
#c("IM" , "IGC" , "DGC") , 
#fill=c(col_risk[3:1]) ,
#horiz=TRUE, cex=0.8)
dev.off()

out_file <- paste0(out_path , "/baseInfo_coef.",molecular_type,".tsv")
write.table( dat , out_file , row.names = F , sep = "\t" , quote = F )